import numpy as np
from matplotlib.pylab import plt

# 一元线性回归模型
class SimpleRegress:
    def __init__(self, remain_dec=2):
        # 斜率
        self.slope = 0.0
        # 截距
        self.intercept = 0.0
        # remain_dec模型默认保留小数位数
        self.remain_dec = remain_dec
        # 临时数据
        self.temp_data = np.array([])

    def fit(self,train_x:list[float], train_y: list[float]) -> None:
        '''
        特征矩阵 feature_matrix

        λ0 = [1,1,1,...,1]
        λ1 = [x1,x2,....xn]
        f = [y1,y2,....yn]

        [(λ0,λ0),(λ0,λ1)] [b]   [(f,λ0)]
        [(λ1,λ0),(λ1,λ1)] [k]   [(f,λ1)]
                 =
                y=kx+b
        '''
        # 留1矩阵
        λ0 = np.ones(len(train_x))
        λ1 = np.array(train_x)
        f = np.array(train_y)
        # 存储临时数据集
        self.temp_data = f;
        λ00, λ01, λ10, λ11 = sum(λ0 * λ0), sum(λ0 * λ1), sum(λ1 * λ0), sum(λ1 * λ1)
        fλ0, fλ1 = sum(f * λ0), sum(f * λ1)
        # 生成斜率
        k = (λ10 * fλ0 - λ00 * fλ1) / (λ01 * λ10 - λ00 * λ11)
        # 生成截距
        b = (λ11 * fλ0 - λ01 * fλ1) / (λ00 * λ11 - λ10 * λ01)
        self.slope = round(k, self.remain_dec)
        self.intercept = round(b,self.remain_dec)

    # 绘制图像
    def drawDataGraph(self,title="semi-liner-regress",xlabel="X-lable",ylabel="Value",color="SeaGreen"):
        # plt.rcParams['font.family'] = 'SimHei'
        data = self.temp_data
        k = self.slope
        b = self.intercept
        plt.figure(figsize=(10, 5), dpi=80)
        plt.title(title)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        ## 一元线性回归方程
        regress_function = np.array([round(k * i + b, 1) for i in range(0, data.size + 2)])
        plt.plot([str(i) for i in range(0, data.size + 2)], regress_function, linewidth=1.5, c=color)
        plt.scatter([i for i in range(0, data.size)], data)
        border = plt.gca()
        border.spines["top"].set_color("none")
        border.spines["right"].set_color("none")
        plt.show()

    # 预测
    def predict(self,x: float)->float:
        return round(self.slope*x+self.intercept,self.remain_dec)
